{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = 'all' # default is ‘last_expr’\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'10.0.0'"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import azure.batch\n",
    "azure.batch.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from azure.batch import BatchServiceClient\n",
    "from azure.batch.models import *\n",
    "from azure.common.credentials import ServicePrincipalCredentials"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set up an instance of the batch processing API\n",
    "\n",
    "We create one Azure Batch Pool for each instance of the batch processing API.\n",
    "\n",
    "The limit for the number of Pools in our Batch account is 100."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Create an Azure Batch Pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MODIFY THIS CELL\n",
    "\n",
    "# POOL_ID should start with the name of the API instance this pool will be used for\n",
    "\n",
    "POOL_ID = 'internal_1'\n",
    "assert len(POOL_ID) <= 64, 'pool_id has more than 64 characters'\n",
    "\n",
    "# choose the account in East US or South Central US\n",
    "BATCH_ACCOUNT_URL = 'https://cameratrapssc.southcentralus.batch.azure.com'  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# secrets read from environment variables\n",
    "REGISTRY_PASSWORD = os.environ['REGISTRY_PASSWORD']\n",
    "STORAGE_ACCOUNT_KEY = os.environ['STORAGE_ACCOUNT_KEY']\n",
    "\n",
    "# authenticate with Batch account using the service principle \"camera-trap-async-api\" in our AAD\n",
    "APP_CLIENT_ID = os.environ['APP_CLIENT_ID']\n",
    "APP_CLIENT_SECRET = os.environ['APP_CLIENT_SECRET']\n",
    "APP_TENANT_ID = os.environ['APP_TENANT_ID']\n",
    "\n",
    "\n",
    "# other configuration info\n",
    "\n",
    "# Docker image\n",
    "REGISTRY_SERVER = 'cameratracrsppftkje.azurecr.io'\n",
    "REGISTRY_USERNAME = REGISTRY_SERVER.split('.')[0]\n",
    "\n",
    "CONTAINER_IMAGE_NAME = 'cameratracrsppftkje.azurecr.io/tensorflow:1.14.0-gpu-py3' # login server/repository:tag\n",
    "\n",
    "# storage\n",
    "STORAGE_ACCOUNT_NAME = 'cameratrap'  # in the engineering subscription\n",
    "\n",
    "# names of two containers supporting the API instances in the above storage account\n",
    "STORAGE_CONTAINER_MODELS = 'models'\n",
    "STORAGE_CONTAINER_API = 'batch-api'\n",
    "\n",
    "# Azure Batch node pool VM type\n",
    "POOL_VM_SIZE = 'Standard_NC6s_v3'  # https://docs.microsoft.com/en-us/azure/virtual-machines/ncv3-series\n",
    "\n",
    "# auto-scale formula - can be set manually in Azure portal\n",
    "# last statement makes sure that nodes aren't removed until their tasks are finished\n",
    "# docs: https://docs.microsoft.com/en-us/azure/batch/batch-automatic-scaling\n",
    "\n",
    "# MODIFY the \"cappedPoolSize\" if it should be other than 16 dedicated nodes\n",
    "POOL_AUTO_SCALE_FORMULA = \"\"\"\n",
    "// In this formula, the pool size is adjusted based on the number of tasks in the queue. \n",
    "// Note that both comments and line breaks are acceptable in formula strings.\n",
    "\n",
    "// Get pending tasks for the past 15 minutes.\n",
    "$samples = $ActiveTasks.GetSamplePercent(TimeInterval_Minute * 15);\n",
    "\n",
    "// If we have fewer than 70 percent data points, we use the last sample point, otherwise we use the maximum of last sample point and the history average.\n",
    "$tasks = $samples < 70 ? max(0, $ActiveTasks.GetSample(1)) : \n",
    "max( $ActiveTasks.GetSample(1), avg($ActiveTasks.GetSample(TimeInterval_Minute * 15)));\n",
    "\n",
    "// If number of pending tasks is not 0, set targetVM to pending tasks, otherwise set to 0, since there is usually long intervals between job submissions.\n",
    "$targetVMs = $tasks > 0 ? $tasks : 0;\n",
    "\n",
    "// The pool size is capped at 16, if target VM value is more than that, set it to 16.\n",
    "cappedPoolSize = 16;\n",
    "$TargetDedicatedNodes = max(0, min($targetVMs, cappedPoolSize));\n",
    "\n",
    "// Set node deallocation mode - keep nodes active only until tasks finish\n",
    "$NodeDeallocationOption = taskcompletion;\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_batch_exception(batch_exception):\n",
    "    \"\"\"\n",
    "    Prints the contents of the specified Batch exception.\n",
    "    \"\"\"\n",
    "    print('-------------------------------------------')\n",
    "    print('Exception encountered:')\n",
    "    if batch_exception.error and \\\n",
    "            batch_exception.error.message and \\\n",
    "            batch_exception.error.message.value:\n",
    "        print(batch_exception.error.message.value)\n",
    "        if batch_exception.error.values:\n",
    "            print()\n",
    "            for msg in batch_exception.error.values:\n",
    "                print(f'{msg.key}:\\t{msg.value}')\n",
    "    print('-------------------------------------------')\n",
    "\n",
    "def create_pool(batch_service_client, pool_id):\n",
    "    \"\"\"\n",
    "    Create a pool with pool_id and the Docker image specified by constants in above cells\n",
    "    \"\"\"\n",
    "    # we have to use VM images supporting GPU access *and* Docker\n",
    "    # this VM image will run our custom container\n",
    "    image_ref = ImageReference(\n",
    "        publisher='microsoft-azure-batch',\n",
    "        offer='ubuntu-server-container',\n",
    "        sku='20-04-lts',\n",
    "        version='latest'  # URN: microsoft-azure-batch:ubuntu-server-container:16-04-lts:1.1.0\n",
    "        # The Azure Batch container image only accepts 'latest' version\n",
    "    )\n",
    "\n",
    "    # specify a container registry from which to pull the custom container\n",
    "    # see the `batch_service` folder on instructions for building the container image\n",
    "    container_registry = ContainerRegistry(\n",
    "        registry_server=REGISTRY_SERVER,\n",
    "        user_name=REGISTRY_USERNAME,\n",
    "        password=REGISTRY_PASSWORD\n",
    "    )\n",
    "\n",
    "    container_conf = ContainerConfiguration(\n",
    "        container_image_names = [CONTAINER_IMAGE_NAME],\n",
    "        container_registries =[container_registry]\n",
    "    )\n",
    "\n",
    "    vm_config = VirtualMachineConfiguration(\n",
    "        image_reference=image_ref,\n",
    "        container_configuration=container_conf,\n",
    "        node_agent_sku_id='batch.node.ubuntu 20.04'\n",
    "    )\n",
    "\n",
    "    # mount the `models` and the `batch-api` blob containers\n",
    "    container_models = MountConfiguration(\n",
    "        azure_blob_file_system_configuration=AzureBlobFileSystemConfiguration(\n",
    "            account_name=STORAGE_ACCOUNT_NAME,\n",
    "            container_name=STORAGE_CONTAINER_MODELS,\n",
    "            relative_mount_path=STORAGE_CONTAINER_MODELS,  # use container name as relative path\n",
    "            account_key=STORAGE_ACCOUNT_KEY,\n",
    "            blobfuse_options='-o attr_timeout=240 -o entry_timeout=240 -o negative_timeout=120 -o allow_other'\n",
    "        )\n",
    "    )\n",
    "    container_batch_api = MountConfiguration(\n",
    "        azure_blob_file_system_configuration=AzureBlobFileSystemConfiguration(\n",
    "            account_name=STORAGE_ACCOUNT_NAME,\n",
    "            container_name=STORAGE_CONTAINER_API,\n",
    "            relative_mount_path=STORAGE_CONTAINER_API,  # use container name as relative path\n",
    "            account_key=STORAGE_ACCOUNT_KEY,\n",
    "            # allow_other needs to be flagged - task running inside container needs to access this blob container\n",
    "            blobfuse_options='-o attr_timeout=240 -o entry_timeout=240 -o negative_timeout=120 -o allow_other'\n",
    "        )\n",
    "    )\n",
    "\n",
    "    new_pool = PoolAddParameter(\n",
    "        id=POOL_ID,\n",
    "        display_name=POOL_ID,\n",
    "\n",
    "        vm_size=POOL_VM_SIZE,\n",
    "        \n",
    "        enable_auto_scale=True,\n",
    "        auto_scale_formula=POOL_AUTO_SCALE_FORMULA,\n",
    "\n",
    "        virtual_machine_configuration=vm_config,\n",
    "\n",
    "        # default is 1; each task occupies the entire GPU so we can only run one task at a time on a node\n",
    "        task_slots_per_node=1,\n",
    "\n",
    "        mount_configuration=[container_models, container_batch_api],\n",
    "    )\n",
    "    batch_service_client.pool.add(new_pool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "credentials = ServicePrincipalCredentials(\n",
    "    client_id=APP_CLIENT_ID,\n",
    "    secret=APP_CLIENT_SECRET,\n",
    "    tenant=APP_TENANT_ID,\n",
    "    resource='https://batch.core.windows.net/'\n",
    ")\n",
    "\n",
    "# if using the Batch quota system, use https://docs.microsoft.com/en-us/python/api/azure-batch/azure.batch.batch_auth.sharedkeycredentials?view=azure-python\n",
    "# to authenticate instead of the service principal is also okay.\n",
    "\n",
    "batch_client = BatchServiceClient(credentials=credentials, batch_url=BATCH_ACCOUNT_URL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 29 ms, sys: 3.54 ms, total: 32.5 ms\n",
      "Wall time: 1.01 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# pool creation should finish in about a minute\n",
    "\n",
    "try:\n",
    "    create_pool(batch_client, POOL_ID)\n",
    "except BatchErrorException as e:\n",
    "    print_batch_exception(e)\n",
    "    raise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Upload the scoring script\n",
    "\n",
    "Note that all instances share this scoring script!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MODIFY THIS CELL\n",
    "\n",
    "# path to the scoring script; modify if cwd is not `api_core`\n",
    "path_scoring_script = 'batch_service/score.py'\n",
    "\n",
    "# SAS with write permission for uploading output JSONs\n",
    "sas_query_str = ''  # get a write-enabled SAS for the container below\n",
    "\n",
    "output_container_url = f'https://cameratrap.blob.core.windows.net/batch-api{sas_query_str}'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# upload the scoring script to the container above; Batch Tasks will retrieve the script from there\n",
    "\n",
    "output_container_client = ContainerClient.from_container_url(output_container_url)\n",
    "\n",
    "with open(path_scoring_script, 'rb') as f:\n",
    "    script_blob_client = output_container_client.upload_blob(name='scripts/score.py', data=f, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Useful CLI commands for using Docker images with Batch\n",
    "\n",
    "List all Batch supported images with their \"capabilities\" (e.g. \"DockerCompatible\", \"NvidiaTeslaDriverInstalled\"):\n",
    "```\n",
    "az batch pool supported-images list\n",
    "```\n",
    "with the pool information provided in additional parameters.\n",
    "\n",
    "Listing all versions of a SKU of image:\n",
    "```\n",
    "az vm image list --all --publisher microsoft-dsvm\n",
    "```\n",
    "\n",
    "You may need to accept the terms of an image:\n",
    "```\n",
    "az vm image list --all --publisher <publisher>\n",
    "```\n",
    "to find the URN for the image you want to use, followed by:\n",
    "\n",
    "```\n",
    "az vm image terms accept --urn <corresponding-urn>\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:cameratraps-batch-api]",
   "language": "python",
   "name": "conda-env-cameratraps-batch-api-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
